import pandas as pd
import numpy as np


class FactorCalculations:
    @staticmethod
    def calculate_momentum_5(data, grouped):
        """5天动量因子 (过去5天收益率)"""
        return grouped['close'].pct_change(5)

    @staticmethod
    def calculate_reversal_20(data, grouped):
        """20天反转因子 (过去20天收益率取反)"""
        return -grouped['close'].pct_change(20)

    @staticmethod
    def calculate_volatility_20(data, grouped):
        """20天波动率因子 (过去20天收益率标准差)"""
        return grouped['close'].pct_change().rolling(20).std()

    @staticmethod
    def calculate_volume_change_5(data, grouped):
        """5天成交量变化率"""
        return grouped['volume'].pct_change(5)

    @staticmethod
    def calculate_pe_ratio(data, _):
        """PE比率 (直接使用数据中的pe列)"""
        return data['pe']

    @staticmethod
    def calculate_pb_ratio(data, _):
        """PB比率 (直接使用数据中的pb列)"""
        return data['pb']

    @staticmethod
    def calculate_size_factor(data, _):
        """市值因子 (市值取对数)"""
        return np.log(data['market_cap'])

    @staticmethod
    def calculate_liquidity(data, _):
        """流动性因子 (成交金额/市值)"""
        return (data['volume'] * data['close']) / data['market_cap']

    @staticmethod
    def calculate_hl_volatility(data, _):
        """高低波动 ((high-low)/close)"""
        return (data['high'] - data['low']) / data['close']

    @staticmethod
    def calculate_rsi_5(data, grouped):
        """RSI近似 (过去5天上涨日比例)"""
        ret = grouped['close'].pct_change()
        return ret.rolling(5).apply(lambda x: (x > 0).mean())